# -*- coding: utf-8 -*-
'''
Created on Jul 21, 2019

@author: yl
'''

from FFT_Interpolation import *
import cv2
from mpl_toolkits.mplot3d import Axes3D
from scipy import signal
from scipy.optimize import curve_fit
from scipy.signal import *

def fit_func(x, a, b, c):
    return a*(x-b)**2 + c

j = complex(0, 1)
c = 3e8 # 光速 [m/s]
Lamda = 633e-9 # 光波长 [m]
Fc = c / Lamda # 光频率 [Hz]
L = 0.1
D_set = np.linspace(0, Lamda*2, 51)

pix_size = 5.3e-6
pix_num = 1280
screen_diameter = pix_num * pix_size
window = signal.gaussian(pix_num, std=pix_num/10)
x, y = np.meshgrid(window, window)
window_2d = x*y
# plt.figure('Gaussian window')
# plt.title("Whole Frame")
# plt.imshow(window_2d, cmap='gray')
# plt.show()
hor_angle_set = []
hor_phase_set = []
num = 0
for D in D_set:
    num = num+1
    I_0 = 110
    dx = np.linspace(0, (pix_num-1)*pix_size, num=pix_num)
    dy = dx
    V_x, V_y, Vz = 0.000, 0.000, 1 # Reference
    V_x_2, V_y_2, Vz_2 = 0.001, 0.000, 1 # Measurement
    tana = (V_x**2 + V_y**2)**0.5
    tana_2 = (V_x_2**2 + V_y_2**2)**0.5
    L = 0.2
    X, Y = np.meshgrid(dx, dy)
    D_d = D - X * np.tan(V_x_2 - V_x) + Y * np.tan(V_y_2 - V_y)
    L_d = L + X * V_x + Y * V_y
    diff_L_d_2 = D_d + (L_d + D_d) * (1 - (V_x_2**2 + V_y_2**2)) / (1 + (V_x_2**2 + V_y_2**2)) - L_d * (1 - (V_x**2 + V_y**2)) / (1 + (V_x**2 + V_y**2))
    Z = I_0 * (1 + np.cos(2 * np.pi * diff_L_d_2 / Lamda))
    
    R = 2
    center = [(pix_num-1+800)*pix_size/2, (pix_num-800)*pix_size/2]
    r = np.sqrt((X-center[0])**2 + (Y-center[1])**2)
    d = R - np.sqrt(R**2 - r**2)
    A = 2 * np.pi * (d+2*D) / Lamda
    I_r = np.sin(A)**2 * I_0/10
    
    R = 2
    center = [(pix_num-1-800)*pix_size/2, (pix_num+800)*pix_size/2]
    r = np.sqrt((X-center[0])**2 + (Y-center[1])**2)
    d = R - np.sqrt(R**2 - r**2)
    A = 2 * np.pi * (d+2*D*0) / Lamda
    I_r_2 = np.sin(A)**2 * I_0/20
    
    Z = (Z+I_r_2)
    Z = Z*window_2d
    
    
    
    hor_sig = Z[640]
    line_noise = np.random.normal(2.43, 0.8, len(hor_sig)) 
    line_noise = line_noise.round().astype(int)
#     hor_sig += line_noise
#     hor_sig = np.diff(hor_sig)
#     hor_sig = np.concatenate(([0], hor_sig))
    DC_num = 800
    hor_freq_estim, hor_phase_estim, hor_freqline, hor_magnitude, hor_phase,  hor_m_k_num, hor_X_m_k, hor_freq_for_phase  = FFT_interpolation_2(hor_sig, pix_size, 1e5, DC_num)
    hor_FFT_start = np.where(hor_magnitude > hor_X_m_k*0.4)[0][0]
    hor_FFT_end = np.where(hor_magnitude > hor_X_m_k*0.4)[0][-1]
    hor_fit_x = hor_freqline[hor_FFT_start:hor_FFT_end+1]
    hor_fit_freq = hor_magnitude[hor_FFT_start:hor_FFT_end+1]
    hor_fit_phase = hor_phase[hor_FFT_start:hor_FFT_end+1]
    
    params = curve_fit(fit_func, hor_fit_x, hor_fit_freq)
    [a, b, c] = params[0]
    hor_f_fit = b
    
    hor_angle = V_x-Lamda*hor_f_fit/2
    hor_angle_set.append(hor_angle)
    print(num, hor_m_k_num, D, hor_phase_estim)
    hor_phase_set.append(hor_phase_estim)
    
#     plt.figure('Simulated Pattern')
#     plt.title("Whole Frame")
#     im = plt.imshow(Z, cmap='gray')
#     plt.colorbar(im, fraction=0.046, pad=0.04)
#     plt.show()
plt.figure('Horizontal line')
plt.plot(hor_sig)
plt.grid(which='major', axis='both')
# plt.show()

plt.figure('Horizontal tilting')
plt.plot((hor_angle_set-np.average(hor_angle_set))*1e6)
plt.grid(which='major', axis='both')
plt.ylabel('urad')

plt.figure('Hor Phase')
plt.plot((np.unwrap(hor_phase_set)/2/np.pi*Lamda/2+D_set)*1e9)
plt.grid(which='major', axis='both')
plt.ylabel('nm')
plt.show()
     
